import cv2
import torch
from Model import Siamese
from Config import Config
import pandas as pd
from torch.utils.data import DataLoader,Dataset
import torch.nn as nn
import albumentations as A
from albumentations.pytorch import ToTensorV2
import csv
import tqdm
test_transformer = A.Compose([
    A.Resize(height=Config.img_size, width=Config.img_size),
    #A.cutout(),
    # A.HorizontalFlip(p=0.5),
    # A.VerticalFlip(p=0.5),
    # A.RandomRotate90(p=0.5),
    # A.RandomCrop(height=224,width=224,p=0.5),
    A.Normalize(  max_pixel_value=255.0, p=1.0),
    ToTensorV2(p=1.0),
])
class testdatset(nn.Module):
    def __init__(self,path,transformer=test_transformer ):
        self.path = path
        self.transformer=transformer
        self.data = pd.read_csv(self.path+"valid_data.csv",sep=',',names=["imageA","imageB"])[1:]
        self.len = len(self.data)
    def __len__(self):
        return self.len
    def __getitem__(self, index):
        index =index+1
        imageA = cv2.imread(self.path+"images/"+self.data["imageA"][index],cv2.COLOR_BGR2RGB)
        imageB = cv2.imread(self.path + "images/" + self.data["imageB"][index], cv2.COLOR_BGR2RGB)
        augementedA = self.transformer(image = imageA)
        augementedB = self.transformer(image=imageB)
        return augementedA['image'],augementedB['image'],self.data["imageA"][index],self.data["imageB"][index]


if __name__ == '__main__':
    model_path = Config.model_dir+Config.model_name+".pth"
    path = "/home/one/lhbdata/pet/pet_biometric_challenge_2022/validation/"
    model = Siamese()

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict['model'])
    model = model.eval()
    file = open(Config.results_dir+'test.csv','w')
    writer = csv.writer(file)
    writer.writerow(['imageA', 'imageB', 'prediction'])
    dataset = testdatset(path)
    dataloader = DataLoader(dataset,batch_size=1,shuffle=False,num_workers=2)
    for i,(img1,img2,img1_name,img2_name) in tqdm(dataloader):
        pred = model(img1,img2)
        #_, preds = torch.max(pred.data, 0)
        #print(pred.item())
        #print(preds)
        #print(img1_name[0])
        writer.writerow([img1_name[0], img2_name[0], pred.item()])
        print(i)


    file.close()



    #data = pd.read_csv(path+"valid_data.csv",sep=',',names=["imageA","imageB"])[1:]

   # print(data["imageA"],data["imageB"])